# from paddlex import create_pipeline

# pipeline = create_pipeline(pipeline="seal_recognition")

# output = pipeline.predict("./test_images/01.png")
# for res in output:
#     rec_texts = list(res.get('ocr_result').get('rec_text'))
#     for i in range(len(rec_texts)):
#         print(f"第{i+1}个印章的内容：",rec_texts[i]) ## 打印预测的结构化输出
#     res.save_to_img("./output_images/") ## 保存可视化结果
# seal_recognition.py
import os
from paddlex import create_pipeline

def seal_reg(image_folder = "./test_images", output_folder = "./output_images"):

    # 创建印章识别管道
    pipeline = create_pipeline(pipeline="seal_recognition")

    # 如果输出文件夹不存在，则创建
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    seals_num = {}
    # 遍历图片文件夹中的所有文件
    for root, dirs, files in os.walk(image_folder):
        for filename in files:
            # 检查文件是否为图片文件
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                # 构建图片文件的完整路径
                image_path = os.path.join(root, filename)
                print(f"正在处理图片: {image_path}")
                try:
                    # 进行预测
                    output = pipeline.predict(image_path)
                    # 获取图片所在的上一级目录名称
                    parent_dir_name = os.path.basename(os.path.dirname(image_path))
                    # 构建输出子文件夹的路径
                    output_subfolder = os.path.join(output_folder, parent_dir_name)
                    # 获取图片所在的上一级目录名称
                    parent_dir_name = os.path.basename(os.path.dirname(image_path))
                    # 构建输出子文件夹的路径
                    output_subfolder = os.path.join(output_folder, parent_dir_name)
                    if parent_dir_name not in seals_num:
                        seals_num[parent_dir_name] = 0
                    for res in output:
                        # 获取预测的印章文本
                        rec_texts = list(res.get('ocr_result').get('rec_text'))
                        seals_num[parent_dir_name] += len(rec_texts)
                        for i in range(len(rec_texts)):
                            print(f"第{i+1}个印章的内容：", rec_texts[i])  # 打印预测的结构化输出
                    
                        # 保存可视化结果到对应的子文件夹中
                        res.save_to_img(output_subfolder)
                        print("成功存储", output_subfolder)
                    
                except Exception as e:
                    print(f"处理图片 {image_path} 时出错: {e}")
    return seals_num